#include "yolox_post_process.h"

#include <sys/time.h>

#include "common.h"
#include "typeConv.hpp"
#include <algorithm>
#include <cmath>
#include <future>
#include <iomanip>
#include <iostream>
#include <string.h>
#include <string>
#include <utility>
#include <vector>
#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

#define CLASS_NUM 80 // 2//80

using namespace COMMON;

struct Yolov5sConfig {
  std::vector<int> strides;
  std::vector<std::vector<std::pair<double, double>>> anchors_table;
  int class_num;
  std::vector<std::string> class_names;
};

static const Yolov5sConfig default_yolov5s_config = {
    {8, 16, 32},
    {{{10, 13}, {16, 30}, {33, 23}},
     {{30, 61}, {62, 45}, {59, 119}},
     {{116, 90}, {156, 198}, {373, 326}}},
    CLASS_NUM,
    {"person",        "bicycle",      "car",
     "motorcycle",    "airplane",     "bus",
     "train",         "truck",        "boat",
     "traffic light", "fire hydrant", "stop sign",
     "parking meter", "bench",        "bird",
     "cat",           "dog",          "horse",
     "sheep",         "cow",          "elephant",
     "bear",          "zebra",        "giraffe",
     "backpack",      "umbrella",     "handbag",
     "tie",           "suitcase",     "frisbee",
     "skis",          "snowboard",    "sports ball",
     "kite",          "baseball bat", "baseball glove",
     "skateboard",    "surfboard",    "tennis racket",
     "bottle",        "wine glass",   "cup",
     "fork",          "knife",        "spoon",
     "bowl",          "banana",       "apple",
     "sandwich",      "orange",       "broccoli",
     "carrot",        "hot dog",      "pizza",
     "donut",         "cake",         "chair",
     "couch",         "potted plant", "bed",
     "dining table",  "toilet",       "tv",
     "laptop",        "mouse",        "remote",
     "keyboard",      "cell phone",   "microwave",
     "oven",          "toaster",      "sink",
     "refrigerator",  "book",         "clock",
     "vase",          "scissors",     "teddy bear",
     "hair drier",    "toothbrush"}};

template <class ForwardIterator>
inline size_t argmin(ForwardIterator first, ForwardIterator last) {
  return std::distance(first, std::min_element(first, last));
}

template <class ForwardIterator>
inline size_t argmax(ForwardIterator first, ForwardIterator last) {
  return std::distance(first, std::max_element(first, last));
}

typedef struct Bbox {
  float xmin;
  float ymin;
  float xmax;
  float ymax;

  Bbox() {}

  Bbox(float xmin, float ymin, float xmax, float ymax)
      : xmin(xmin), ymin(ymin), xmax(xmax), ymax(ymax) {}

  ~Bbox() {}
} Bbox;

typedef struct Detection {
  int id;
  float score;
  Bbox bbox;
  const char *class_name = nullptr;
  Detection() {}

  Detection(int id, float score, Bbox bbox)
      : id(id), score(score), bbox(bbox) {}

  Detection(int id, float score, Bbox bbox, const char *class_name)
      : id(id), score(score), bbox(bbox), class_name(class_name) {}

  friend bool operator>(const Detection &lhs, const Detection &rhs) {
    return (lhs.score > rhs.score);
  }

  ~Detection() {}
} Detection;

void yolo5Nms(std::vector<Detection> &input, float iou_threshold, int top_k,
              std::vector<Detection> &result, bool suppress) {
  std::stable_sort(input.begin(), input.end(), std::greater<Detection>());

  std::vector<bool> skip(input.size(), false);

  std::vector<float> areas;
  areas.reserve(input.size());
  for (size_t i = 0; i < input.size(); i++) {
    float width = input[i].bbox.xmax - input[i].bbox.xmin;
    float height = input[i].bbox.ymax - input[i].bbox.ymin;
    areas.push_back(width * height);
  }

  int count = 0;
  for (size_t i = 0; /*count < top_k && */ i < skip.size(); i++) {
    if (skip[i]) {
      continue;
    }
    skip[i] = true;
    ++count;

    for (size_t j = i + 1; j < skip.size(); ++j) {
      if (skip[j]) {
        continue;
      }
      if (suppress == false) {
        if (input[i].id != input[j].id) {
          continue;
        }
      }

      float xx1 = std::max(input[i].bbox.xmin, input[j].bbox.xmin);
      float yy1 = std::max(input[i].bbox.ymin, input[j].bbox.ymin);
      float xx2 = std::min(input[i].bbox.xmax, input[j].bbox.xmax);
      float yy2 = std::min(input[i].bbox.ymax, input[j].bbox.ymax);

      if (xx2 > xx1 && yy2 > yy1) {
        float area_intersection = (xx2 - xx1) * (yy2 - yy1);
        float iou_ratio =
            area_intersection / (areas[j] + areas[i] - area_intersection);
        if (iou_ratio > iou_threshold) {
          skip[j] = true;
        }
      }
    }
    result.push_back(input[i]);
  }
}

void tensorpostProcess(const Yolov5sConfig &yolov5s_config, void *tensor, YoloxPostProcessInfo_t *post_info,
                       int layer, std::vector<Detection> &dets) {
  // auto *data = reinterpret_cast<uint16_t *>(tensor);
  void *data = tensor;
  int class_num = yolov5s_config.class_num;
  int stride = yolov5s_config.strides[layer];
  int num_pred = yolov5s_config.class_num + 4 + 1;

  std::vector<int16_t> class_pred(yolov5s_config.class_num, 0);

  double h_ratio = post_info->height * 1.0 / post_info->ori_height;
  double w_ratio = post_info->width * 1.0 / post_info->ori_width;
  double resize_ratio = std::min(w_ratio, h_ratio);
  if (post_info->is_pad_resize) {
    w_ratio = resize_ratio;
    h_ratio = resize_ratio;
  }

  int grid_height, grid_width;
  grid_height = post_info->height / stride;
  grid_width = post_info->width / stride;

  int16_t box_score_threshold = float2half(post_info->score_threshold);
  for (int h = 0; h < grid_height; h++) {
    for (int w = 0; w < grid_width; w++) {
      for (size_t k = 0; k < static_cast<size_t>(post_info->anchorSize); k++) {
        int16_t *cur_data = (int16_t *)data + k * num_pred;

        int16_t objness = cur_data[4];
        if (objness < box_score_threshold /*post_info->score_threshold*/) {
          continue;
        }

#if 0
int16_t id = cur_data[5];
double confidence = half2float(objness) * half2float(cur_data[6]);
#else
        for (int index = 0; index < class_num; ++index) {
          class_pred[index] = (cur_data[5 + index]);
        }

        int16_t id = argmax(class_pred.begin(), class_pred.end());

        double confidence = half2float(objness) * half2float(class_pred[id]);
#endif

        if (confidence < post_info->score_threshold) {
          continue;
        }

#if 1
        float center_x = half2float(cur_data[0]);
        float center_y = half2float(cur_data[1]);
        float scale_x = half2float(cur_data[2]);
        float scale_y = half2float(cur_data[3]);
#else
        float center_x = cur_data[0] * post_info->width;
        float center_y = cur_data[1] * post_info->height;
        float scale_x = cur_data[2] * post_info->width;
        float scale_y = cur_data[3] * post_info->height;
#endif

        double xmin = (center_x - scale_x / 2.0);
        double ymin = (center_y - scale_y / 2.0);
        double xmax = (center_x + scale_x / 2.0);
        double ymax = (center_y + scale_y / 2.0);
        double w_padding =
            (post_info->width - w_ratio * post_info->ori_width) / 2.0;
        double h_padding =
            (post_info->height - h_ratio * post_info->ori_height) / 2.0;

        double xmin_org = (xmin - w_padding) / w_ratio;
        double xmax_org = (xmax - w_padding) / w_ratio;
        double ymin_org = (ymin - h_padding) / h_ratio;
        double ymax_org = (ymax - h_padding) / h_ratio;

        if (xmax_org <= 0 || ymax_org <= 0) {
          continue;
        }

        if (xmin_org > xmax_org || ymin_org > ymax_org) {
          continue;
        }

        xmin_org = std::max(xmin_org, 0.0);
        xmax_org = std::min(xmax_org, post_info->ori_width - 1.0);
        ymin_org = std::max(ymin_org, 0.0);
        ymax_org = std::min(ymax_org, post_info->ori_height - 1.0);

        Bbox bbox(xmin_org, ymin_org, xmax_org, ymax_org);
        dets.push_back(
            Detection((int)id, confidence, bbox,
                      yolov5s_config.class_names[(int)id].c_str()));
      }
      data = (int16_t *)data + num_pred * post_info->anchorSize;
    }
  }
}

int lynYoloxPostProcess(YoloxPostProcessInfo_t *post_info) {
  int16_t *output_tensor = nullptr;
  lynBoxesInfo *boxesInfo = nullptr;

#ifdef LYNXI_PLUGIN
  output_tensor = (int16_t *)lynPluginGetVirtAddr(post_info->output_tensor);
  if (output_tensor == nullptr) {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  boxesInfo = (lynBoxesInfo *)lynPluginGetVirtAddr(post_info->boxesInfo);
  if (boxesInfo == nullptr) {
    LOG_PLUGIN_E("get boxesInfo addr error\n");
    return -1;
  }
#else
  output_tensor = (int16_t *)post_info->output_tensor;
  boxesInfo = post_info->boxesInfo;
#endif

  Yolov5sConfig yolov5s_config = default_yolov5s_config;
  generateClassNameAndNum(post_info->labelList,
                          yolov5s_config.class_names,
                          yolov5s_config.class_num);
  void *tensor0 = NULL;
  tensor0 = output_tensor;
  void *tensor1 = NULL;
  int grid_width1 = post_info->width / yolov5s_config.strides[0];
  int grid_height1 = post_info->height / yolov5s_config.strides[0];
  tensor1 = (int16_t *)tensor0 + 1 * post_info->anchorSize * grid_width1 *
                                     grid_height1 *
                                     (yolov5s_config.class_num + 5);
  void *tensor2 = NULL;
  int grid_width2 = post_info->width / yolov5s_config.strides[1];
  int grid_height2 = post_info->height / yolov5s_config.strides[1];
  tensor2 = (int16_t *)tensor1 + 1 * post_info->anchorSize * grid_width2 *
                                     grid_height2 *
                                     (yolov5s_config.class_num + 5);
  std::vector<Detection> dets;
  std::vector<Detection> det_restuls;

  tensorpostProcess(yolov5s_config, tensor0, post_info, 0, dets);
  tensorpostProcess(yolov5s_config, tensor1, post_info, 1, dets);
  tensorpostProcess(yolov5s_config, tensor2, post_info, 2, dets);

  yolo5Nms(dets, post_info->nms_threshold, post_info->nms_top_k, det_restuls,
           false);

  boxesInfo->boxesNum = std::min((int)det_restuls.size(), BOX_MAX_NUM);

  for (int i = 0; i < boxesInfo->boxesNum; i++) {
    boxesInfo->boxes[i].xmax = det_restuls[i].bbox.xmax;
    boxesInfo->boxes[i].xmin = det_restuls[i].bbox.xmin;
    boxesInfo->boxes[i].ymax = det_restuls[i].bbox.ymax;
    boxesInfo->boxes[i].ymin = det_restuls[i].bbox.ymin;
    boxesInfo->boxes[i].score = det_restuls[i].score;
    boxesInfo->boxes[i].id = det_restuls[i].id;
    strncpy(boxesInfo->boxes[i].label, det_restuls[i].class_name,
            sizeof(boxesInfo->boxes[i].label) - 1);
    if (static_cast<int>(boxesInfo->boxes[i].xmax) > post_info->ori_width - 1)
      boxesInfo->boxes[i].xmax = post_info->ori_width - 1;
    if (static_cast<int>(boxesInfo->boxes[i].xmin) < 0)
      boxesInfo->boxes[i].xmin = 0;
    if (static_cast<int>(boxesInfo->boxes[i].ymax) > post_info->ori_height - 1)
      boxesInfo->boxes[i].ymax = post_info->ori_height - 1;
    if (static_cast<int>(boxesInfo->boxes[i].ymin) < 0)
      boxesInfo->boxes[i].ymin = 0;
  }

  return 0;
}
